"""
HIT-UAV YOLO 标签划分脚本 - 增强版
根据 normal_json 中的 train/val/test 划分，组织 yolo_labels
"""

import os
import shutil
from pathlib import Path

# ============= 配置路径 =============
BASE_DIR = r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\HIT-UAV\HIT-UAV"

# 源路径
YOLO_LABELS_DIR = os.path.join(BASE_DIR, r"HIT-UAV-download\yolo_labels")
NORMAL_JSON_DIR = os.path.join(BASE_DIR, "normal_json")

# 目标路径（新建文件夹存放划分后的标签）
OUTPUT_DIR = os.path.join(BASE_DIR, "yolo_labels_split")

# ============= 主函数 =============
def split_yolo_labels():
    """根据 normal_json 的划分，组织 yolo_labels"""

    print("=" * 70)
    print("HIT-UAV YOLO 标签划分工具 v2.0")
    print("=" * 70)

    # 检查路径是否存在
    if not os.path.exists(YOLO_LABELS_DIR):
        print(f"❌ 错误：找不到 YOLO 标签目录：{YOLO_LABELS_DIR}")
        return

    if not os.path.exists(NORMAL_JSON_DIR):
        print(f"❌ 错误：找不到参考目录：{NORMAL_JSON_DIR}")
        return

    # 创建输出目录
    print(f"\n📁 创建输出目录: {OUTPUT_DIR}")
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # 统计信息
    stats = {
        'train': {'images': 0, 'labels_found': 0, 'labels_missing': 0},
        'val': {'images': 0, 'labels_found': 0, 'labels_missing': 0},
        'test': {'images': 0, 'labels_found': 0, 'labels_missing': 0}
    }

    # 处理每个数据集划分
    for split in ['train', 'val', 'test']:
        print(f"\n{'='*70}")
        print(f"📁 处理 {split.upper()} 集")
        print(f"{'='*70}")

        # 获取该划分下的所有图片文件名
        split_img_dir = os.path.join(NORMAL_JSON_DIR, split)

        if not os.path.exists(split_img_dir):
            print(f"⚠️  跳过：目录不存在 {split_img_dir}")
            continue

        # 创建输出子目录
        output_split_dir = os.path.join(OUTPUT_DIR, split)
        os.makedirs(output_split_dir, exist_ok=True)
        print(f"   输出目录: {output_split_dir}")

        # 获取所有图片文件
        image_files = []
        for ext in ['.jpg', '.jpeg', '.png', '.bmp', '.JPG', '.JPEG', '.PNG', '.BMP']:
            image_files.extend(Path(split_img_dir).glob(f'*{ext}'))

        stats[split]['images'] = len(image_files)
        print(f"   发现图片: {len(image_files)} 张")

        # 复制对应的标签文件
        missing_labels = []

        for idx, img_path in enumerate(image_files, 1):
            # 显示进度
            if idx % 500 == 0 or idx == len(image_files):
                print(f"   进度: {idx}/{len(image_files)}", end='\r')

            # 获取图片文件名（不含扩展名）
            img_stem = img_path.stem

            # 对应的标签文件路径
            label_filename = f"{img_stem}.txt"
            src_label_path = os.path.join(YOLO_LABELS_DIR, label_filename)
            dst_label_path = os.path.join(output_split_dir, label_filename)

            # 复制标签文件
            if os.path.exists(src_label_path):
                shutil.copy2(src_label_path, dst_label_path)
                stats[split]['labels_found'] += 1
            else:
                missing_labels.append(label_filename)
                stats[split]['labels_missing'] += 1

        print()  # 换行
        print(f"   ✅ 成功复制: {stats[split]['labels_found']} 个标签")
        print(f"   ⚠️  缺失标签: {stats[split]['labels_missing']} 个")

        # 显示部分缺失的标签（如果有）
        if missing_labels and len(missing_labels) <= 10:
            print(f"   缺失的标签文件:")
            for label in missing_labels[:10]:
                print(f"      - {label}")
        elif missing_labels:
            print(f"   前10个缺失的标签:")
            for label in missing_labels[:10]:
                print(f"      - {label}")
            print(f"      ... 还有 {len(missing_labels)-10} 个")

    # 输出统计信息
    print(f"\n{'='*70}")
    print("📊 最终统计")
    print(f"{'='*70}")

    total_images = sum(s['images'] for s in stats.values())
    total_found = sum(s['labels_found'] for s in stats.values())
    total_missing = sum(s['labels_missing'] for s in stats.values())

    print(f"{'数据集':<10} {'图片数':<10} {'标签数':<10} {'缺失':<10} {'覆盖率':<10}")
    print("-" * 70)

    for split in ['train', 'val', 'test']:
        images = stats[split]['images']
        found = stats[split]['labels_found']
        missing = stats[split]['labels_missing']
        coverage = f"{found/images*100:.1f}%" if images > 0 else "N/A"

        print(f"{split:<10} {images:<10} {found:<10} {missing:<10} {coverage:<10}")

    coverage_total = f"{total_found/total_images*100:.1f}%" if total_images > 0 else "N/A"
    print("-" * 70)
    print(f"{'总计':<10} {total_images:<10} {total_found:<10} {total_missing:<10} {coverage_total:<10}")

    print(f"\n✅ 标签文件已保存到: {OUTPUT_DIR}")

    # 生成目录结构说明
    print(f"\n📂 生成的目录结构:")
    print(f"{OUTPUT_DIR}/")
    print(f"├── train/  ({stats['train']['labels_found']} 个标签)")
    print(f"├── val/    ({stats['val']['labels_found']} 个标签)")
    print(f"└── test/   ({stats['test']['labels_found']} 个标签)")

    return stats


# ============= 创建 YOLO 数据集配置文件 =============
def create_yaml_config():
    """创建 YOLO 训练配置文件"""

    yaml_content = f"""# HIT-UAV 数据集配置文件
# 用于 YOLOv8 训练

# 数据集路径
path: {BASE_DIR.replace(os.sep, '/')}  # 数据集根目录

# 图片路径
train: normal_json/train
val: normal_json/val
test: normal_json/test

# 注意：YOLO 会自动在图片同目录下查找 labels/ 文件夹
# 如果标签不在图片同目录，需要手动指定：
# train: 
#   images: normal_json/train
#   labels: yolo_labels_split/train

# 类别数量
nc: 4

# 类别名称
names:
  0: Person
  1: Bicycle
  2: Car
  3: OtherVehicle

# 训练参数建议
# imgsz: 640        # 图像大小
# batch: 16         # 批次大小
# epochs: 100       # 训练轮数
# device: 0         # GPU 设备
# workers: 8        # 数据加载线程数
# project: runs/train
# name: hit_uav
"""

    yaml_path = os.path.join(OUTPUT_DIR, "hit_uav_dataset.yaml")
    with open(yaml_path, 'w', encoding='utf-8') as f:
        f.write(yaml_content)

    print(f"\n📄 已生成 YOLO 配置文件: {yaml_path}")
    print("\n💡 训练命令示例:")
    print(f"   yolo train data=\"{yaml_path}\" model=yolov8n.pt epochs=100")

    return yaml_path


# ============= 可选：将标签移动到图片目录 =============
def copy_labels_to_image_dirs():
    """
    将标签文件复制到图片目录下的 labels/ 子文件夹
    这是 YOLO 的标准目录结构
    """

    print(f"\n{'='*70}")
    print("📋 可选操作：将标签复制到图片目录")
    print(f"{'='*70}")

    user_input = input("是否将标签复制到图片目录下的 labels/ 文件夹？(y/n): ")

    if user_input.lower() != 'y':
        print("跳过此步骤")
        return

    for split in ['train', 'val', 'test']:
        # 源标签目录
        src_label_dir = os.path.join(OUTPUT_DIR, split)

        # 目标标签目录（图片目录下的 labels/）
        img_dir = os.path.join(NORMAL_JSON_DIR, split)
        dst_label_dir = os.path.join(img_dir, 'labels')

        if os.path.exists(src_label_dir):
            os.makedirs(dst_label_dir, exist_ok=True)

            label_files = list(Path(src_label_dir).glob('*.txt'))
            print(f"\n复制 {split} 集标签: {len(label_files)} 个文件")

            for label_file in label_files:
                dst_path = os.path.join(dst_label_dir, label_file.name)
                shutil.copy2(label_file, dst_path)

            print(f"   ✅ 已复制到: {dst_label_dir}")


# ============= 执行 =============
if __name__ == "__main__":
    try:
        # 执行标签划分
        stats = split_yolo_labels()

        # 生成配置文件
        yaml_path = create_yaml_config()

        # 询问是否复制到图片目录
        copy_labels_to_image_dirs()

        print("\n" + "=" * 70)
        print("✅ 所有操作完成！")
        print("=" * 70)

        print("\n📌 下一步:")
        print("   1. 检查生成的文件夹是否正确")
        print("   2. 使用生成的 YAML 文件开始训练")
        print("   3. 如果标签覆盖率低，检查是否有空白图片")

    except Exception as e:
        print(f"\n❌ 错误：{e}")
        import traceback
        traceback.print_exc()

    input("\n按回车键退出...")